import torch
import torch.nn as nn
import torch.nn.functional as F



class SuperpixelAggregation(nn.Module):
    def __init__(self, height, width, grid_size=(12, 12), num_iters=5):
        """
        初始化 SPA 模块
        :param height: 特征图高度 H
        :param width: 特征图宽度 W
        :param grid_size: 超像素初始化网格大小 (Hs, Ws)
        :param num_iters: EM 聚类迭代次数
        """
        super(SuperpixelAggregation, self).__init__()
        self.H, self.W = height, width
        self.Hs, self.Ws = grid_size
        self.num_iters = num_iters

    def forward(self, x):
        """
        :param x: 输入特征图 [B, C, H, W]
        :return:
            s: 聚类后的超像素特征 [B, M, C]
            A: 最终关联图 [B, N, M]，表示每个像素属于每个超像素的概率
        """
        B, C, H, W = x.shape
        N = H * W  # 总像素数
        M_H = H // self.Hs
        M_W = W // self.Ws
        M = M_H * M_W  # 总超像素数

        # ========== Step 1: Patchify - 初始化超像素特征 s0 ==========
        # 平均池化作为初始中心，尺寸为 M_H x M_W
        s = F.adaptive_avg_pool2d(x, output_size=(M_H, M_W))  # [B, C, M_H, M_W]
        s = s.view(B, C, -1).permute(0, 2, 1).contiguous()     # [B, M, C]

        # 展开 x 为像素序列
        x_flat = x.view(B, C, -1).permute(0, 2, 1).contiguous()  # [B, N, C]

        for _ in range(self.num_iters):
            # ========== Step 2a: 计算距离（欧氏距离平方）并转为权重 ==========
            # dist[i, j] = ||x[i] - s[j]||^2
            dist = torch.cdist(x_flat, s, p=2) ** 2  # [B, N, M]
            A = torch.exp(-dist)                    # soft assignment 权重 [B, N, M]

            # ========== Step 2b: 用 A 对 x 进行加权平均，更新超像素 s ==========
            z = A.sum(dim=1, keepdim=True) + 1e-8   # 防止除以0 [B, 1, M]
            s = torch.bmm(A.transpose(1, 2), x_flat) / z.transpose(1, 2)  # [B, M, C]

        return s, A  # s: [B, M, C], A: [B, N, M]
    



if __name__=='__main__':
    picture=torch.rand(size=[8, 24, 144, 144], dtype=torch.float32)
    model=SuperpixelAggregation(height=144, width=144)
    s, a=model(picture)
    print(s.shape, a.shape)